// UDP Sokect 管理
// Created by kenny on 25-10-20.
//
#include "mcp_udp.h"
#include "mcp_logger.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#ifdef _WIN32
    #include <ws2tcpip.h>
    #if defined(_MSC_VER)
       #pragma comment(lib, "ws2_32.lib")
    #endif
#else
    #include <sys/socket.h>
    #include <netinet/in.h>
    #include <arpa/inet.h>
    #include <unistd.h>
    #include <pthread.h>
#endif

// 内部函数声明
#ifdef _WIN32
static DWORD WINAPI udp_recv_thread_func(LPVOID param);
#else
static void* udp_recv_thread_func(void* param);
#endif

int mcp_udp_init(mcp_udp_t* manager, const char* ip, int port,
                           udp_recv_callback_t recv_cb) {
    if (!manager || !ip || !recv_cb) {
        return -1;
    }

    memset(manager, 0, sizeof(mcp_udp_t));

    strncpy(manager->server_ip, ip, sizeof(manager->server_ip) - 1);
    manager->server_port = port;
    manager->recv_callback = recv_cb;

    manager->socket_fd = INVALID_SOCKET;
    manager->is_connected = false;
    manager->running = false;

    // 初始化Winsock（Windows平台）
#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        LOG_ERROR("WSAStartup failed");
        return -1;
    }
#endif

    return 0;
}

#ifdef _WIN32
static DWORD WINAPI udp_recv_thread_func(LPVOID param) {
#else
static void* udp_recv_thread_func(void* param) {
#endif
    mcp_udp_t* manager = (mcp_udp_t*)param;
    uint8_t buffer[1024];

    struct sockaddr_in client_addr;
    socklen_t client_addr_len = sizeof(client_addr);

    while (manager->running) {
        int received;
#ifdef _WIN32
        received = recvfrom(manager->socket_fd, (char*)buffer, sizeof(buffer), 0,
                           (struct sockaddr*)&client_addr, &client_addr_len);
#else
        received = recvfrom(manager->socket_fd, buffer, sizeof(buffer), 0,
                           (struct sockaddr*)&client_addr, &client_addr_len);
#endif

        if (received == SOCKET_ERROR) {
#ifdef _WIN32
            if (WSAGetLastError() == WSAETIMEDOUT) {
                continue; // 超时，继续接收
            }
#endif
            LOG_ERROR("UDP receive failed");
            continue;
        }

        if (received > 0 && manager->recv_callback) {
            char client_ip[16];
            inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, sizeof(client_ip));
            int client_port = ntohs(client_addr.sin_port);

            manager->recv_callback(buffer, received, client_ip, client_port);
        }
    }

#ifdef _WIN32
    return 0;
#else
    return NULL;
#endif
}

int mcp_udp_connect(mcp_udp_t* manager) {
    if (!manager) return -1;

    struct sockaddr_in server_addr;

    // 创建UDP socket
    manager->socket_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (manager->socket_fd == INVALID_SOCKET) {
        LOG_ERROR("UDP socket creation failed");
        return -1;
    }

    // 设置服务器地址（用于默认发送目标）
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(manager->server_port);
    server_addr.sin_addr.s_addr = inet_addr(manager->server_ip);

    // UDP不需要连接操作，但我们可以设置默认目标地址
    // 这里只是标记为"已连接"状态
    manager->is_connected = true;
    manager->running = true;

    // 设置接收超时
#ifdef _WIN32
    DWORD timeout = 5000; // 5秒超时
    setsockopt(manager->socket_fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&timeout, sizeof(timeout));
#else
    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;
    setsockopt(manager->socket_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
#endif

    LOG_INFO("UDP socket created for %s:%d", manager->server_ip, manager->server_port);

    // 启动接收线程
#ifdef _WIN32
    manager->recv_thread = CreateThread(NULL, 0, udp_recv_thread_func, manager, 0, NULL);
    if (manager->recv_thread == NULL) {
        LOG_ERROR("Failed to create UDP receive thread");
        mcp_udp_close(manager);
        return -1;
    }
#else
    if (pthread_create(&manager->recv_thread, NULL, udp_recv_thread_func, manager) != 0) {
        LOG_ERROR("Failed to create UDP receive thread");
        mcp_udp_close(manager);
        return -1;
    }
#endif

    return 0;
}

int mcp_udp_send(mcp_udp_t* manager, const uint8_t* data, uint16_t length) {
    if (!manager || !data || length == 0) {
        return -1;
    }

    if (!manager->is_connected) {
        LOG_WARN("Cannot send data, UDP socket not ready");
        return -1;
    }

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(manager->server_port);
    server_addr.sin_addr.s_addr = inet_addr(manager->server_ip);

    int sent;
#ifdef _WIN32
    sent = sendto(manager->socket_fd, (const char*)data, length, 0,
                 (const struct sockaddr*)&server_addr, sizeof(server_addr));
#else
    sent = sendto(manager->socket_fd, data, length, 0,
                 (const struct sockaddr*)&server_addr, sizeof(server_addr));
#endif

    if (sent == SOCKET_ERROR) {
        LOG_ERROR("UDP send failed");
        return -1;
    }

    LOG_DEBUG("Sent %d bytes via UDP to %s:%d", sent, manager->server_ip, manager->server_port);
    return 0;
}

int mcp_udp_send_to(mcp_udp_t* manager, const uint8_t* data, uint16_t length,
                              const char* dest_ip, int dest_port) {
    if (!manager || !data || length == 0 || !dest_ip) {
        return -1;
    }

    if (!manager->is_connected) {
        LOG_WARN("Cannot send data, UDP socket not ready");
        return -1;
    }

    struct sockaddr_in dest_addr;
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port);
    dest_addr.sin_addr.s_addr = inet_addr(dest_ip);

    int sent;
#ifdef _WIN32
    sent = sendto(manager->socket_fd, (const char*)data, length, 0,
                 (const struct sockaddr*)&dest_addr, sizeof(dest_addr));
#else
    sent = sendto(manager->socket_fd, data, length, 0,
                 (const struct sockaddr*)&dest_addr, sizeof(dest_addr));
#endif

    if (sent == SOCKET_ERROR) {
        LOG_ERROR("UDP send to %s:%d failed", dest_ip, dest_port);
        return -1;
    }

    LOG_DEBUG("Sent %d bytes via UDP to %s:%d", sent, dest_ip, dest_port);
    return 0;
}

int mcp_udp_close(mcp_udp_t* manager) {
    if (!manager) return -1;

    manager->running = false;
    manager->is_connected = false;

    if (manager->socket_fd != INVALID_SOCKET) {
#ifdef _WIN32
        closesocket(manager->socket_fd);
#else
        close(manager->socket_fd);
#endif
        manager->socket_fd = INVALID_SOCKET;
    }

    // 等待接收线程结束
#ifdef _WIN32
    if (manager->recv_thread) {
        WaitForSingleObject(manager->recv_thread, INFINITE);
        CloseHandle(manager->recv_thread);
        manager->recv_thread = NULL;
    }
#else
    if (manager->recv_thread) {
        pthread_join(manager->recv_thread, NULL);
        manager->recv_thread = 0;
    }
#endif

    LOG_INFO("UDP socket closed");
    return 0;
}

bool mcp_udp_is_connected(mcp_udp_t* manager) {
    return manager && manager->is_connected;
}

void mcp_udp_cleanup(mcp_udp_t* manager) {
    if (!manager) return;

    mcp_udp_close(manager);

#ifdef _WIN32
    WSACleanup();
#endif
}
